package org.example.springaitest.controller;

import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestController
public class SseController {

    // 存储所有订阅者的SseEmitter（键为客户端ID，值为SseEmitter）
    private static final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();

    // 客户端订阅通知的接口（无超时设置）
    @GetMapping("/sse/subscribe")
    public SseEmitter subscribe() {
        // 生成唯一客户端ID（实际项目中可使用UUID或用户ID）
        String clientId = String.valueOf(System.currentTimeMillis());
        SseEmitter emitter = new SseEmitter(); // 不设置超时时间（默认无超时）

        // 存储emitter，用于后续发送通知
        emitters.put(clientId, emitter);

        // 设置回调处理（仅保留完成、错误回调，移除超时回调）
        emitter.onCompletion(() -> emitters.remove(clientId));
        emitter.onError(e -> emitters.remove(clientId));

        // 发送初始化消息，确认连接成功
        try {
            emitter.send(SseEmitter.event()
                    .id("init")
                    .name("connected")
                    .data("已连接，等待工具调用通知..."));
        } catch (IOException e) {
            emitter.completeWithError(e);
        }

        return emitter;
    }

    // 静态方法：用于工具调用时发送通知
    public static void sendToolInvokeNotify(String toolName) {
        String message = String.format(
                "工具【%s】已被调用",
                toolName
        );

        // 遍历所有订阅者，发送通知
        emitters.forEach((clientId, emitter) -> {
            try {
                emitter.send(SseEmitter.event()
                        .id(clientId)
                        .name("toolInvoke")
                        .data(message));
            } catch (IOException e) {
                // 发送失败则移除emitter
                emitter.completeWithError(e);
                emitters.remove(clientId);
            }
        });
    }
}